package com.dtsheep.elasticsearch.es.client;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchScrollRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * ScrollSearchUtil
 *
 * <p>
 *     滚动查询
 * </p>
 *
 * @author chuanbai
 * @date 2024/6/12
 */
public class ScrollSearchUtil {


    @Resource
    private RestHighLevelClient client;
    @Resource
    private ObjectMapper objectMapper;
    /**
     * 滚动查询
     */
    public <T> List<T> rollSearch(Class<T> clazz, String indices) throws IOException {

        SearchRequest searchRequest = new SearchRequest(indices);

        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.query(QueryBuilders.matchAllQuery());
        searchSourceBuilder.size(100); // 设置每次返回的文档数量
        searchRequest.scroll(TimeValue.timeValueMinutes(1)); // 设置滚动时间窗口为1分钟
        searchRequest.source(searchSourceBuilder);

        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        List<T> documents = new ArrayList<>();

        while (true) {
            for (SearchHit hit : searchResponse.getHits().getHits()) {
                // 将每个搜索命中转换为指定类型的对象并添加到列表中
                T document = objectMapper.readValue(hit.getSourceAsString(), clazz);
                documents.add(document);
            }

            // 检查是否还有数据需要检索
            String scrollId = searchResponse.getScrollId();
            SearchScrollRequest scrollRequest = new SearchScrollRequest(scrollId);
            scrollRequest.scroll(TimeValue.timeValueMinutes(1));
            searchResponse = client.scroll(scrollRequest, RequestOptions.DEFAULT);

            if (searchResponse.getHits().getHits().length == 0) {
                // 所有数据都已经检索完毕
                break;
            }
        }

        return documents;
    }

}
